package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"
)

var (
	flagDir = flag.String("dir", "images", "Directory to scan for images")
	// flagLang  = flag.String("lang", "chi_sim+eng", "Tesseract language(s)")
	flagLang  = flag.String("lang", "chi_sim", "Tesseract language(s)")
	flagJobs  = flag.Int("jobs", 4, "Number of concurrent OCR workers")
	flagQuiet = flag.Bool("q", false, "Quiet mode (less logs)")
)

// DateTimeInfo 存储提取的时间信息
type DateTimeInfo struct {
	PaymentDate  string
	PaymentTime  string
	TransferDate string
	TransferTime string
	OriginalName string
	NewName      string
}

// 预编译日期识别正则（包含常见OCR误识别）
var datePatterns = []*regexp.Regexp{
	regexp.MustCompile(`(转账时间|转账肘间)\s*(\d{4})年(\d{1,2})月(\d{1,2})[日曰El]`),
	regexp.MustCompile(`(支付时间)\s*(\d{4})年(\d{1,2})月(\d{1,2})[日曰El]`),
	regexp.MustCompile(`(创建时间|付款时间|发货时间)[：:]?\s*(\d{4})[\-_/](\d{1,2})[\-_/](\d{1,2})`),
	regexp.MustCompile(`(支付时间|转账时间)[：:]?\s*(\d{4})[\-_/](\d{1,2})[\-_/](\d{1,2})`),
	// 无标签的纯中文日期（作为兜底）
	regexp.MustCompile(`()(\d{4})年(\d{1,2})月(\d{1,2})[日曰El]`),
}

// 预编译时间识别正则
var (
	timeColonPattern    = regexp.MustCompile(`\b(\d{1,2}):(\d{1,2})(?::(\d{1,2}))?\b`)
	timeCompactPattern  = regexp.MustCompile(`\b(\d{6})\b`)
	alreadyNamedPattern = regexp.MustCompile(`^\d{4}-\d{2}-\d{2}(?:_\d{6})?\.[a-zA-Z0-9]+$`)
)

// 中文数字到阿拉伯数字的映射
var chineseNumbers = map[string]string{
	"一": "1", "二": "2", "三": "3", "四": "4", "五": "5",
	"六": "6", "七": "7", "八": "8", "九": "9", "十": "10",
	"零": "0", "〇": "0",
}

// 月份中文映射
var chineseMonths = map[string]string{
	"一月": "01", "二月": "02", "三月": "03", "四月": "04",
	"五月": "05", "六月": "06", "七月": "07", "八月": "08",
	"九月": "09", "十月": "10", "十一月": "11", "十二月": "12",
	"1月": "01", "2月": "02", "3月": "03", "4月": "04",
	"5月": "05", "6月": "06", "7月": "07", "8月": "08",
	"9月": "09", "10月": "10", "11月": "11", "12月": "12",
}

func main() {
	// 添加全局panic恢复
	defer func() {
		if r := recover(); r != nil {
			fmt.Printf("程序发生严重错误，但尝试继续运行: %v\n", r)
		}
	}()

	flag.Parse()

	// 收集需要处理的图片文件
	dir := *flagDir
	var files []string
	var imageCount int

	if !*flagQuiet {
		fmt.Printf("扫描目录: %s\n", dir)
	}

	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			if !*flagQuiet {
				fmt.Printf("访问文件时出错: %s (错误: %v)\n", path, err)
			}
			return nil // 继续处理其他文件
		}

		// 跳过目录
		if info.IsDir() {
			return nil
		}

		// 检查是否为图片文件
		if isImageFile(path) {
			imageCount++
			// 检查是否已经重命名过
			base := filepath.Base(path)
			if !alreadyNamedPattern.MatchString(base) {
				files = append(files, path)
			} else if !*flagQuiet {
				fmt.Printf("跳过已命名文件: %s\n", base)
			}
		}

		return nil
	})

	if err != nil {
		log.Fatal(err)
	}

	if len(files) == 0 {
		if !*flagQuiet {
			fmt.Printf("未找到需要处理的图片文件\n")
		}
		return
	}

	if !*flagQuiet {
		fmt.Printf("找到 %d 个图片文件，需要处理 %d 个\n", imageCount, len(files))
	}

	// 并发处理
	jobs := *flagJobs
	if jobs < 1 {
		jobs = 1
	}

	// 创建信号量控制并发数
	sem := make(chan struct{}, jobs)
	var wg sync.WaitGroup
	var processedCount int32
	var mu sync.Mutex

	startTime := time.Now()

	for _, f := range files {
		f := f
		wg.Add(1)
		sem <- struct{}{}

		go func() {
			defer wg.Done()
			defer func() { <-sem }()

			// 添加panic恢复，确保单个文件错误不会影响整个程序
			defer func() {
				if r := recover(); r != nil {
					if !*flagQuiet {
						fmt.Printf("处理文件时发生panic，跳过: %s (错误: %v)\n", f, r)
					}
				}
			}()

			// 处理文件
			processOneFile(f, *flagLang, *flagQuiet)

			// 更新计数
			mu.Lock()
			processedCount++
			current := processedCount
			mu.Unlock()

			if !*flagQuiet && current%10 == 0 {
				fmt.Printf("已处理: %d/%d 文件\n", current, len(files))
			}
		}()
	}

	// 等待所有goroutine完成
	wg.Wait()

	elapsed := time.Since(startTime)

	if !*flagQuiet {
		fmt.Printf("\n处理完成！\n")
		fmt.Printf("总图片文件: %d 个\n", imageCount)
		fmt.Printf("成功处理: %d 个\n", processedCount)
		fmt.Printf("处理耗时: %v\n", elapsed)
		fmt.Printf("平均速度: %.2f 文件/秒\n", float64(processedCount)/elapsed.Seconds())
	}
}

// 支持的图片文件扩展名
var imageExtensions = map[string]bool{
	".jpg":  true,
	".jpeg": true,
	".png":  true,
	".bmp":  true,
	".tiff": true,
	".tif":  true,
	".gif":  true,
	".webp": true,
	".ico":  true,
}

// 检查是否为图片文件
func isImageFile(path string) bool {
	// 检查文件扩展名
	ext := strings.ToLower(filepath.Ext(path))
	if !imageExtensions[ext] {
		return false
	}

	// 检查文件是否存在且不是目录
	info, err := os.Stat(path)
	if err != nil || info.IsDir() {
		return false
	}

	// 检查文件大小（至少要有内容）
	if info.Size() == 0 {
		return false
	}

	return true
}

// 处理单个文件
func processOneFile(path, lang string, quiet bool) {
	if !quiet {
		fmt.Printf("处理文件: %s\n", path)
	}

	// 提取OCR文本
	text, err := extractTextFromImageWithLang(path, lang)
	if err != nil {
		log.Printf("OCR提取失败 %s: %v", path, err)
		return
	}

	// 解析时间信息
	dateTimeInfo := parseDateTimeInfo(text, path)

	// 生成新文件名
	newName := generateNewFileName(dateTimeInfo)

	// 重命名文件
	if newName != "" {
		if err := renameFile(path, newName); err != nil {
			log.Printf("重命名失败 %s: %v", path, err)
		} else if !quiet {
			fmt.Printf("重命名成功: %s -> %s\n", filepath.Base(path), newName)
		}
	} else if !quiet {
		fmt.Printf("未找到有效时间信息: %s\n", path)
	}
}

// 从图片提取OCR文本
func extractTextFromImageWithLang(imagePath, lang string) (string, error) {
	// 如果没有指定语言，使用默认值
	if lang == "" {
		lang = "chi_sim+eng"
	}

	cmd := exec.Command("tesseract", imagePath, "stdout", "-l", lang)
	output, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("tesseract执行失败: %v", err)
	}
	return string(output), nil
}

// 解析时间信息
func parseDateTimeInfo(text, filePath string) DateTimeInfo {
	info := DateTimeInfo{
		OriginalName: filepath.Base(filePath),
	}

	// 清理文本
	text = strings.ReplaceAll(text, " ", "")
	text = strings.ReplaceAll(text, "\n", "")
	text = strings.ReplaceAll(text, "\r", "")
	// 日期锚定 + 邻近时间窗口提取
	for _, re := range datePatterns {
		loc := re.FindStringSubmatchIndex(text)
		if loc == nil {
			continue
		}
		// 分组: 1=标签 2=年 3=月 4=日
		subs := re.FindStringSubmatch(text)
		if len(subs) < 5 {
			continue
		}
		label := subs[1]
		year := subs[2]
		month := fmt.Sprintf("%02s", subs[3])
		day := fmt.Sprintf("%02s", subs[4])
		dateStr := fmt.Sprintf("%s-%s-%s", year, month, day)

		// 在日期后固定窗口中找时间（扩大窗口以覆盖多行案例）
		windowStart := loc[1]
		windowEnd := windowStart + 120
		if windowEnd > len(text) {
			windowEnd = len(text)
		}
		window := normalizeForTime(text[windowStart:windowEnd])

		if hhmmss, ok := extractValidTime(window); ok {
			if strings.Contains(label, "支付") {
				info.PaymentDate = dateStr
				info.PaymentTime = hhmmss
			} else {
				info.TransferDate = dateStr
				info.TransferTime = hhmmss
			}
			// 找到一组就足够
			break
		} else {
			// 即使未找到时间，也先记录日期
			if strings.Contains(label, "支付") {
				if info.PaymentDate == "" {
					info.PaymentDate = dateStr
				}
			} else {
				if info.TransferDate == "" {
					info.TransferDate = dateStr
				}
			}
		}
	}

	return info
}

// 规范化时间窗口内容，修正常见OCR错误
func normalizeForTime(s string) string {
	r := s
	r = strings.ReplaceAll(r, "二", ":")
	r = strings.ReplaceAll(r, "：", ":")
	r = strings.ReplaceAll(r, ";", ":")
	r = strings.ReplaceAll(r, "·", ":")
	r = strings.ReplaceAll(r, "O", "0")
	r = strings.ReplaceAll(r, "o", "0")
	r = strings.ReplaceAll(r, "l", "1")
	r = strings.ReplaceAll(r, "I", "1")
	r = strings.ReplaceAll(r, "|", "1")
	r = strings.ReplaceAll(r, "S", "5")
	r = strings.ReplaceAll(r, "s", "5")
	return r
}

// 从窗口中提取并校验时间
func extractValidTime(window string) (string, bool) {
	// 优先匹配带冒号格式 H:M:S 或 H:M
	if m := timeColonPattern.FindStringSubmatch(window); len(m) >= 3 {
		hh, mm, ss := m[1], m[2], "00"
		if len(m) == 4 && m[3] != "" {
			ss = m[3]
		}
		if v, ok := validateHHMMSS(hh, mm, ss); ok {
			return v, true
		}
	}

	// 其次匹配纯数字紧凑格式 HHMMSS
	if m := timeCompactPattern.FindStringSubmatch(window); len(m) == 2 {
		s := m[1]
		hh, mm, ss := s[0:2], s[2:4], s[4:6]
		if v, ok := validateHHMMSS(hh, mm, ss); ok {
			return v, true
		}
	}
	return "", false
}

func validateHHMMSS(hh, mm, ss string) (string, bool) {
	// 简单范围校验
	toInt := func(x string) int {
		n := 0
		for _, ch := range x {
			if ch < '0' || ch > '9' {
				return 100
			}
			n = n*10 + int(ch-'0')
		}
		return n
	}
	h := toInt(hh)
	m := toInt(mm)
	s := toInt(ss)
	if h < 0 || h > 23 || m < 0 || m > 59 || s < 0 || s > 59 {
		return "", false
	}
	return fmt.Sprintf("%02s:%02s:%02s", hh, mm, ss), true
}

// 生成新文件名
func generateNewFileName(info DateTimeInfo) string {
	var dateStr, timeStr string

	// 优先使用支付时间，如果没有则使用转账时间
	if info.PaymentDate != "" {
		dateStr = info.PaymentDate
		timeStr = info.PaymentTime
	} else if info.TransferDate != "" {
		dateStr = info.TransferDate
		timeStr = info.TransferTime
	}

	if dateStr == "" {
		return ""
	}

	// 格式化时间部分
	if timeStr == "" {
		timeStr = "000000"
	} else {
		timeStr = strings.ReplaceAll(timeStr, ":", "")
	}

	// 生成新文件名
	ext := filepath.Ext(info.OriginalName)
	newName := fmt.Sprintf("%s_%s%s", dateStr, timeStr, ext)

	return newName
}

// 重命名文件
func renameFile(oldPath, newName string) error {
	dir := filepath.Dir(oldPath)
	newPath := filepath.Join(dir, newName)

	// 检查新文件是否已存在
	if _, err := os.Stat(newPath); err == nil {
		// 文件已存在，添加序号
		ext := filepath.Ext(newName)
		baseName := strings.TrimSuffix(newName, ext)
		counter := 1
		for {
			newNameWithCounter := fmt.Sprintf("%s_%d%s", baseName, counter, ext)
			newPath = filepath.Join(dir, newNameWithCounter)
			if _, err := os.Stat(newPath); os.IsNotExist(err) {
				break
			}
			counter++
		}
	}

	return os.Rename(oldPath, newPath)
}
